from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


def make_cuda_ext(name, sources):

    return CUDAExtension(name='{}'.format(name),
                         sources=[p for p in sources],
                         extra_compile_args={
                             'cxx': [],
                             'nvcc': [
                                 '-D__CUDA_NO_HALF_OPERATORS__',
                                 '-D__CUDA_NO_HALF_CONVERSIONS__',
                                 '-D__CUDA_NO_HALF2_OPERATORS__',
                             ]
                         })


setup(name='deform_conv',
      ext_modules=[
          make_cuda_ext(name='deform_conv_cuda', sources=['src/deform_conv_cuda.cpp', 'src/deform_conv_cuda_kernel.cu'])
      ],
      cmdclass={'build_ext': BuildExtension},
      zip_safe=False)
